import pandas as pd
import os

def check_stock_pattern(df, 
                       up_days=5,         # 上升段总天数
                       up_percent=0.7,    # 上升段内上涨天数百分比(70%)
                       up_gain=0.2,       # 上升段总涨幅(20%)
                       down_days=5,       # 下跌段总天数
                       down_percent=0.7,  # 下跌段内下跌天数百分比(70%)
                       down_loss=0.15,    # 下跌段总跌幅(15%)
                       stable_days=3,     # 平稳段总天数
                       stable_volatility=0.05,  # 平稳段价格波动率(5%)
                       stable_volume_threshold=1e6):  # 平稳段成交量阈值
    """
    检查股票是否符合"上升-回落-平稳"形态
    
    参数:
    df: 股票数据DataFrame
    up_days: 上升段总天数
    up_percent: 上升段内上涨天数百分比(0-1)
    up_gain: 上升段总涨幅(0-1)，即最终价格/起始价格 - 1
    down_days: 下跌段总天数
    down_percent: 下跌段内下跌天数百分比(0-1)
    down_loss: 下跌段总跌幅(0-1)，即(峰值价格-最终价格)/峰值价格
    stable_days: 平稳段总天数
    stable_volatility: 平稳段价格波动率(0-1)
    stable_volume_threshold: 平稳段成交量阈值
    """
    # 按日期升序排序
    df = df.sort_values(by='日期')
    close_prices = df['收盘价'].values
    volumes = df['成交量'].values
    n = len(close_prices)
    
    # 确保有足够的数据点
    if n < up_days + down_days + stable_days:
        return False
    
    # 1. 寻找上升段
    required_up_days = int(up_days * up_percent)  # 至少需要的上涨天数
    up_end = -1
    
    # 滑动窗口检查上升段
    for i in range(up_days, n - down_days - stable_days + 1):
        window = close_prices[i-up_days:i]
        # 计算上涨天数
        up_count = sum(window[j] > window[j-1] for j in range(1, up_days))
        # 计算总涨幅
        gain = (window[-1] - window[0]) / window[0]
        
        if up_count >= required_up_days and gain >= up_gain:
            up_end = i  # 记录上升段结束位置
            break
    
    if up_end == -1:
        return False
    
    # 2. 寻找下跌段
    required_down_days = int(down_days * down_percent)  # 至少需要的下跌天数
    peak_price = close_prices[up_end - 1]  # 上升段结束时的价格作为峰值
    down_end = -1
    
    # 从上升段结束后开始检查下跌段
    for j in range(up_end + down_days, n - stable_days + 1):
        window = close_prices[j-down_days:j]
        # 计算下跌天数
        down_count = sum(window[k] < window[k-1] for k in range(1, down_days))
        # 计算总跌幅
        loss = (peak_price - window[-1]) / peak_price
        
        if down_count >= required_down_days and loss >= down_loss:
            down_end = j  # 记录下跌段结束位置
            break
    
    if down_end == -1:
        return False
    
    # 3. 寻找平稳段
    # 从下跌段结束后开始检查平稳段
    if down_end + stable_days > n:
        return False
    
    stable_window = close_prices[down_end:down_end+stable_days]
    # 计算价格波动率
    stable_vol = (stable_window.max() - stable_window.min()) / stable_window.min()
    # 计算平均成交量
    volume_avg = volumes[down_end:down_end+stable_days].mean()
    
    # 检查是否满足平稳条件
    if stable_vol <= stable_volatility and volume_avg <= stable_volume_threshold:
        return True
    
    return False


def process_csv_files(folder_path, **kwargs):
    """
    处理文件夹中的所有CSV文件，寻找符合形态的股票
    
    参数:
    folder_path: CSV文件所在文件夹路径
   ** kwargs: 传递给check_stock_pattern的参数
    """
    result = []
    for file_name in os.listdir(folder_path):
        if file_name.endswith('.csv'):
            file_path = os.path.join(folder_path, file_name)
            try:
                df = pd.read_csv(file_path, encoding='utf-8')
                # 按股票代码分组处理
                grouped = df.groupby('代码')
                for code, group in grouped:
                    if check_stock_pattern(group, **kwargs):
                        result.append({
                            '代码': code,
                            '股票名称': group['股票名称'].iloc[0],
                            '文件': file_name
                        })
            except Exception as e:
                print(f"处理文件 {file_name} 时出错: {str(e)}")
    return result


if __name__ == '__main__':
    # 替换为实际的CSV文件路径
    folder_path = 'your_csv_folder_path'
    
    # 设置参数，可以根据需要调整
    params = {
        'up_days': 5,         # 上升段5天
        'up_percent': 0.7,    # 70%天数上涨
        'up_gain': 0.2,       # 总涨幅20%
        'down_days': 5,       # 下跌段5天
        'down_percent': 0.7,  # 70%天数下跌
        'down_loss': 0.15,    # 总跌幅15%
        'stable_days': 3,     # 平稳段3天
        'stable_volatility': 0.05,  # 波动率5%以内
        'stable_volume_threshold': 1e6  # 成交量阈值
    }
    
    matched_stocks = process_csv_files(folder_path,** params)
    
    if matched_stocks:
        print("找到符合形态的股票：")
        for i, stock in enumerate(matched_stocks, 1):
            print(f"{i}. {stock['股票名称']}（{stock['代码']}），来自文件：{stock['文件']}")
    else:
        print("未找到符合形态的股票")
